use crate::common::*;
use crate::common_types::*;
use crate::traits::{FragmentShader, Program, VertexShader};
use glam::{ivec2, ivec4, mat4, vec2, vec4, Mat4, Vec2, Vec4};
use std::ffi::CStr;
use std::mem;
use super::strcmp;

//
// C++ Structs translated to Rust
//

#[derive(Clone, Copy, Debug, Default)]
struct QuadSegment {
    rect: RectWithEndpoint,
    uv_rect: RectWithEndpoint,
}

#[derive(Clone, Copy, Debug, Default)]
struct QuadPrimitive {
    bounds: RectWithEndpoint,
    clip: RectWithEndpoint,
    uv_rect: RectWithEndpoint,
    pattern_scale_offset: vec4,
    color: vec4,
}

#[derive(Clone, Copy, Debug, Default)]
struct QuadHeader {
    transform_id: i32,
    z_id: i32,
    pattern_input: ivec2,
}

#[derive(Clone, Copy, Debug, Default)]
struct QuadInstance {
    prim_address_i: i32,
    prim_address_f: i32,
    quad_flags: i32,
    edge_flags: i32,
    part_index: i32,
    segment_index: i32,
    picture_task_address: i32,
}

#[derive(Clone, Copy, Debug, Default)]
struct VertexInfo {
    local_pos: vec2,
}

#[derive(Clone, Copy, Debug, Default)]
struct PrimitiveInfo {
    local_pos: vec2,
    local_prim_rect: RectWithEndpoint,
    local_clip_rect: RectWithEndpoint,
    segment: QuadSegment,
    edge_flags: i32,
    quad_flags: i32,
    pattern_input: ivec2,
}

//
// Common Struct
//

#[derive(Clone, Debug, Default)]
struct PsQuadTexturedCommon {
    attrib_locations: AttribLocations,
    u_transform: mat4,
    // Flat varyings
    v_color: vec4,
    v_flags: ivec4,
    v_uv0_sample_bounds: vec4,
}

#[derive(Clone, Debug, Default)]
struct AttribLocations {
    a_position: usize,
    a_data: usize,
}

impl AttribLocations {
    fn bind_loc(&mut self, name: &CStr, index: i32) {
        if strcmp(name, "aPosition") {
            self.a_position = index as usize;
        } else if strcmp(name, "aData") {
            self.a_data = index as usize;
        }
    }

    fn get_loc(&self, name: &CStr) -> i32 {
        if strcmp(name, "aPosition") {
            if self.a_position != NULL_ATTRIB { self.a_position as i32 } else { -1 }
        } else if strcmp(name, "aData") {
            if self.a_data != NULL_ATTRIB { self.a_data as i32 } else { -1 }
        } else {
            -1
        }
    }
}

//
// Vertex Shader
//

#[derive(Clone, Debug, Default)]
struct PsQuadTexturedVert {
    common: PsQuadTexturedCommon,
    // Inputs
    a_position: vec2,
    a_data: ivec4,
    // Outputs
    v_uv0: vec2,
    gl_position: vec4,
}

#[repr(C)]
#[derive(Default, Copy, Clone)]
struct InterpOutputs {
    v_uv0: vec2,
}

impl PsQuadTexturedVert {
    fn main(&mut self, context: &mut ShaderContext) {
        let prim = self.quad_primive_info(context);
        if (prim.quad_flags & 16) != 0 {
            self.common.v_flags.z = 1;
        } else {
            self.common.v_flags.z = 0;
        }
        self.antialiasing_vertex(prim);
        self.pattern_vertex(context, prim);
    }
    
    fn decode_instance(&self) -> QuadInstance {
        QuadInstance {
            prim_address_i: self.a_data.x,
            prim_address_f: self.a_data.y,
            quad_flags: (self.a_data.z >> 24) & 255,
            edge_flags: (self.a_data.z >> 16) & 255,
            part_index: (self.a_data.z >> 8) & 255,
            segment_index: self.a_data.z & 255,
            picture_task_address: self.a_data.w,
        }
    }

    fn fetch_header(&self, address: i32, context: &ShaderContext) -> QuadHeader {
        let header = context.fetch_from_gpu_buffer_1i(address);
        QuadHeader {
            transform_id: header.x,
            z_id: header.y,
            pattern_input: header.zw(),
        }
    }

    fn fetch_primitive(&self, index: i32, context: &ShaderContext) -> QuadPrimitive {
        let texels = context.fetch_from_gpu_buffer_5f(index);
        QuadPrimitive {
            bounds: RectWithEndpoint { p0: texels[0].xy(), p1: texels[0].zw() },
            clip: RectWithEndpoint { p0: texels[1].xy(), p1: texels[1].zw() },
            uv_rect: RectWithEndpoint { p0: texels[2].xy(), p1: texels[2].zw() },
            pattern_scale_offset: texels[3],
            color: texels[4],
        }
    }

    fn fetch_segment(&self, base: i32, index: i32, context: &ShaderContext) -> QuadSegment {
        let texels = context.fetch_from_gpu_buffer_2f(base + 5 + (index * 2));
        QuadSegment {
            rect: RectWithEndpoint { p0: texels[0].xy(), p1: texels[0].zw() },
            uv_rect: RectWithEndpoint { p0: texels[1].xy(), p1: texels[1].zw() },
        }
    }

    fn edge_aa_offset(&self, edge: i32, flags: i32) -> f32 {
        if (flags & edge) != 0 { 2.0 } else { 0.0 }
    }

    fn write_vertex(&mut self, local_pos: Vec2, z: f32, transform: &Transform, content_origin: Vec2, task_rect: RectWithEndpoint, device_pixel_scale: f32, quad_flags: i32) -> VertexInfo {
        let mut vi = VertexInfo::default();
        let world_pos = transform.m * local_pos.extend(0.0).extend(1.0);
        let mut device_pos = world_pos.xy() * device_pixel_scale;

        if (quad_flags & 2) != 0 {
            let device_clip_rect = RectWithEndpoint { p0: content_origin, p1: content_origin + (task_rect.p1 - task_rect.p0) };
            device_pos = device_pos.clamp(device_clip_rect.p0, device_clip_rect.p1);
            vi.local_pos = (transform.inv_m * (device_pos / device_pixel_scale).extend(0.0).extend(1.0)).xy();
        } else {
            vi.local_pos = local_pos;
        }

        let final_offset = -content_origin + task_rect.p0;
        self.gl_position = self.common.u_transform * vec4(
            device_pos.x + (final_offset.x * world_pos.w),
            device_pos.y + (final_offset.y * world_pos.w),
            z * world_pos.w,
            world_pos.w,
        );
        vi
    }

    fn scale_offset_map_point(&self, scale_offset: Vec4, p: Vec2) -> Vec2 {
        p * scale_offset.xy() + scale_offset.zw()
    }
    
    fn scale_offset_map_rect(&self, scale_offset: Vec4, r: RectWithEndpoint) -> RectWithEndpoint {
        RectWithEndpoint {
            p0: self.scale_offset_map_point(scale_offset, r.p0),
            p1: self.scale_offset_map_point(scale_offset, r.p1),
        }
    }

    fn quad_primive_info(&mut self, context: &mut ShaderContext) -> PrimitiveInfo {
        let qi = self.decode_instance();
        let qh = self.fetch_header(qi.prim_address_i, context);
        let transform = context.fetch_transform(qh.transform_id);
        let task = context.fetch_picture_task(qi.picture_task_address);
        let prim = self.fetch_primitive(qi.prim_address_f, context);
        let z = qh.z_id as f32;

        let mut seg = if qi.segment_index == 255 {
            QuadSegment { rect: prim.bounds, uv_rect: prim.uv_rect }
        } else {
            self.fetch_segment(qi.prim_address_f, qi.segment_index, context)
        };
        
        let mut local_coverage_rect = seg.rect;
        local_coverage_rect.p0 = local_coverage_rect.p0.max(prim.clip.p0);
        local_coverage_rect.p1 = local_coverage_rect.p1.min(prim.clip.p1);
        local_coverage_rect.p1 = local_coverage_rect.p1.max(local_coverage_rect.p0);

        // swgl_antiAlias calls would go here, but they modify external state.
        match qi.part_index {
            1 => local_coverage_rect.p1.x = local_coverage_rect.p0.x + 2.0,
            2 => {
                local_coverage_rect.p0.x += 2.0;
                local_coverage_rect.p1.x -= 2.0;
                local_coverage_rect.p1.y = local_coverage_rect.p0.y + 2.0;
            },
            3 => local_coverage_rect.p0.x = local_coverage_rect.p1.x - 2.0,
            4 => {
                local_coverage_rect.p0.x += 2.0;
                local_coverage_rect.p1.x -= 2.0;
                local_coverage_rect.p0.y = local_coverage_rect.p1.y - 2.0;
            },
            0 => {
                local_coverage_rect.p0.x += self.edge_aa_offset(1, qi.edge_flags);
                local_coverage_rect.p1.x -= self.edge_aa_offset(4, qi.edge_flags);
                local_coverage_rect.p0.y += self.edge_aa_offset(2, qi.edge_flags);
                local_coverage_rect.p1.y -= self.edge_aa_offset(8, qi.edge_flags);
            },
            _ => {},
        }
        
        let local_pos = local_coverage_rect.p0.lerp(local_coverage_rect.p1, self.a_position);
        
        let mut device_pixel_scale = task.device_pixel_scale;
        if (qi.quad_flags & 4) != 0 {
            device_pixel_scale = 1.0;
        }

        let vi = self.write_vertex(local_pos, z, &transform, task.content_origin, task.task_rect, device_pixel_scale, qi.quad_flags);
        
        self.common.v_color = prim.color;
        let pattern_tx = prim.pattern_scale_offset;
        seg.rect = self.scale_offset_map_rect(pattern_tx, seg.rect);

        PrimitiveInfo {
            local_pos: self.scale_offset_map_point(pattern_tx, vi.local_pos),
            local_prim_rect: self.scale_offset_map_rect(pattern_tx, prim.bounds),
            local_clip_rect: self.scale_offset_map_rect(pattern_tx, prim.clip),
            segment: seg,
            edge_flags: qi.edge_flags,
            quad_flags: qi.quad_flags,
            pattern_input: qh.pattern_input,
        }
    }
    
    fn antialiasing_vertex(&mut self, _prim: PrimitiveInfo) {
        // No-op in this shader
    }
    
    fn pattern_vertex(&mut self, context: &ShaderContext, info: PrimitiveInfo) {
        if info.segment.uv_rect.p0 != info.segment.uv_rect.p1 {
            self.common.v_flags.x = 1;
            self.common.v_color = vec4::ONE;
            let f = (info.local_pos - info.segment.rect.p0) / context.rect_size(info.segment.rect);
            self.vs_init_sample_color0(context, f, info.segment.uv_rect);
        } else {
            self.common.v_flags.x = 0;
        }
    }

    fn vs_init_sample_color0(&mut self, context: &ShaderContext, sample_pos: Vec2, uv_rect: RectWithEndpoint) {
        let uv = uv_rect.p0.lerp(uv_rect.p1, sample_pos);
        let texture_size = context.texture_size(SamplerId::SColor0, 0);
        self.v_uv0 = uv / texture_size;
        self.common.v_uv0_sample_bounds = vec4(
            uv_rect.p0.x + 0.5, uv_rect.p0.y + 0.5,
            uv_rect.p1.x - 0.5, uv_rect.p1.y - 0.5,
        ) / texture_size.extend(texture_size).xyxy();
    }
}

impl VertexShader for PsQuadTexturedVert {
    fn init_batch(&mut self) {}

    fn load_attribs(&mut self, attribs: &[&VertexAttrib], start: u32, instance: i32, _count: i32) {
        unsafe {
            let a_pos_attrib = &*attribs[self.common.attrib_locations.a_position];
            let a_data_attrib = &*attribs[self.common.attrib_locations.a_data];
            let pos_ptr = (a_pos_attrib.data as *const u8).add(a_pos_attrib.stride * start as usize) as *const Vec2;
            self.a_position = *pos_ptr;
            let data_ptr = (a_data_attrib.data as *const u8).add(a_data_attrib.stride * instance as usize) as *const ivec4;
            self.a_data = *data_ptr;
        }
    }

    fn run_primitive(&mut self, context: &mut ShaderContext, interps: *mut u8, interp_stride: usize) {
        self.main(context);
        unsafe {
            let mut dest_ptr = interps as *mut InterpOutputs;
            for _ in 0..4 {
                (*dest_ptr).v_uv0 = self.v_uv0;
                dest_ptr = (dest_ptr as *mut u8).add(interp_stride) as *mut InterpOutputs;
            }
        }
    }

    fn set_uniform_1i(&mut self, _index: i32, _value: i32) {}
    fn set_uniform_4fv(&mut self, _index: i32, _value: &[f32; 4]) {}
    fn set_uniform_matrix4fv(&mut self, index: i32, value: &[f32; 16]) {
        if index == 5 { self.common.u_transform = Mat4::from_cols_array(value); }
    }
}

//
// Fragment Shader
//

#[derive(Clone, Debug, Default)]
struct PsQuadTexturedFrag {
    vert: PsQuadTexturedVert,
    // Varying inputs
    v_uv0: vec2,
    interp_step: InterpOutputs,
    interp_perspective: InterpOutputs,
}

impl PsQuadTexturedFrag {
    fn antialiasing_fragment(&self) -> f32 { 1.0 }

    fn fs_sample_color0(&self, context: &ShaderContext) -> vec4 {
        let uv = self.v_uv0.clamp(self.vert.common.v_uv0_sample_bounds.xy(), self.vert.common.v_uv0_sample_bounds.zw());
        context.texture(SamplerId::SColor0, uv)
    }

    fn pattern_fragment(&self, context: &ShaderContext, mut color: vec4) -> vec4 {
        if self.vert.common.v_flags.x != 0 {
            let texel = self.fs_sample_color0(context);
            color *= texel;
        }
        color
    }

    fn main(&self, context: &ShaderContext) -> vec4 {
        let mut base_color = self.vert.common.v_color;
        base_color *= self.antialiasing_fragment();
        
        let mut output_color = self.pattern_fragment(context, base_color);
        
        if self.vert.common.v_flags.z != 0 {
            output_color = vec4::splat(output_color.r);
        }
        
        output_color
    }
}

impl FragmentShader for PsQuadTexturedFrag {
    fn read_interp_inputs(&mut self, init: *const u8, step: *const u8) {
        unsafe {
            let init = &*(init as *const InterpOutputs);
            let step = &*(step as *const InterpOutputs);
            self.v_uv0 = init.v_uv0;
            self.interp_step.v_uv0 = step.v_uv0 * 4.0;
        }
    }

    fn read_perspective_inputs(&mut self, init: *const u8, step: *const u8, w: f32) {
        unsafe {
            let init = &*(init as *const InterpOutputs);
            let step = &*(step as *const InterpOutputs);
            let inv_w = 1.0 / w;
            self.interp_perspective.v_uv0 = init.v_uv0;
            self.v_uv0 = self.interp_perspective.v_uv0 * inv_w;
            self.interp_step.v_uv0 = step.v_uv0 * 4.0;
        }
    }

    fn run(&mut self, context: &mut ShaderContext) {
        let color = self.main(context);
        context.write_output(color);
        self.step_interp_inputs(4);
    }
    
    fn skip(&mut self, steps: i32) {
        self.step_interp_inputs(steps);
    }

    fn run_perspective(&mut self, context: &mut ShaderContext, next_w: &[f32; 4]) {
        let color = self.main(context);
        context.write_output(color);
        self.step_perspective_inputs(4, next_w);
    }

    fn skip_perspective(&mut self, steps: i32, next_w: &[f32; 4]) {
        self.step_perspective_inputs(steps, next_w);
    }

    fn draw_span_rgba8(&mut self, context: &mut ShaderContext) -> i32 {
        if self.vert.common.v_flags.x != 0 {
            if self.vert.common.v_flags.z != 0 {
                // This path is empty in C++, suggesting it's not a fast path.
                return 0;
            } else {
                context.commit_texture_linear_color_rgba8(SamplerId::SColor0, self.v_uv0, self.vert.common.v_uv0_sample_bounds, self.vert.common.v_color);
            }
        } else {
            context.commit_solid_rgba8(self.vert.common.v_color);
        }
        1
    }
}

impl PsQuadTexturedFrag {
    fn step_interp_inputs(&mut self, steps: i32) {
        let chunks = steps as f32 * 0.25;
        self.v_uv0 += self.interp_step.v_uv0 * chunks;
    }

    fn step_perspective_inputs(&mut self, steps: i32, next_w: &[f32; 4]) {
        let chunks = steps as f32 * 0.25;
        let inv_w = 1.0 / next_w[0];
        self.interp_perspective.v_uv0 += self.interp_step.v_uv0 * chunks;
        self.v_uv0 = self.interp_perspective.v_uv0 * inv_w;
    }
}

//
// Program
//

#[derive(Clone, Debug, Default)]
pub struct PsQuadTexturedProgram {
    frag: PsQuadTexturedFrag,
}

pub fn loader() -> Box<dyn Program> {
    Box::new(PsQuadTexturedProgram::default())
}

impl Program for PsQuadTexturedProgram {
    fn get_vertex_shader(&mut self) -> &mut dyn VertexShader {
        &mut self.frag.vert
    }

    fn get_fragment_shader(&mut self) -> &mut dyn FragmentShader {
        &mut self.frag
    }
    
    fn get_uniform(&self, name: &CStr) -> i32 {
        if strcmp(name, "sColor0") { return 6; }
        if strcmp(name, "sGpuBufferF") { return 3; }
        if strcmp(name, "sGpuBufferI") { return 4; }
        if strcmp(name, "sRenderTasks") { return 2; }
        if strcmp(name, "sTransformPalette") { return 1; }
        if strcmp(name, "uTransform") { return 5; }
        -1
    }
    
    fn get_attrib(&self, name: &CStr) -> i32 {
        self.frag.vert.common.attrib_locations.get_loc(name)
    }

    fn bind_attrib(&mut self, name: &CStr, index: i32) {
        self.frag.vert.common.attrib_locations.bind_loc(name, index);
    }
    
    fn interpolants_size(&self) -> usize {
        mem::size_of::<InterpOutputs>()
    }
}
